from langchain_core.embeddings import Embeddings
from langchain_ollama import OllamaEmbeddings
from langchain_community.vectorstores import Milvus
from langgraph_customer.utils.config import Settings


class VectorStoreRetriever:
    """
    向量检索器,通过工厂方法直接创建milvus和embedding向量模型
    """
    def __init__(self):
        self.embedder: Embeddings = OllamaEmbeddings(
            model=Settings.EMBEDDING_MODEL,
            base_url=Settings.EMBEDDING_URL
        )
        self.vectorstore = Milvus(
            embedding_function=self.embedder,
            connection_args={
                "host": Settings.MILVUS_HOST,
                "port": Settings.MILVUS_PORT,
            },
            collection_name=Settings.MILVUS_COLLECTION_NAME,
        )
        self.retriever = self.vectorstore.as_retriever(
            search_type="similarity",
            search_kwargs={"k": Settings.RETRIEVE_NUM},
        )

    def put_documents(self, docs):
        return self.vectorstore.add_documents(docs)

    def get_similar_docs(self, query) -> list:
        return self.retriever.get_relevant_documents(query)